import tensorflow as tf
import numpy as np
import pandas as pd
from tensorflow.keras import models, layers
import struct, json

# 数据路径
TRAIN_IMG_PATH = 'MNIST_data/train-images-idx3-ubyte/train-images.idx3-ubyte'
TRAIN_LABEL_PATH = 'MNIST_data/train-labels-idx1-ubyte/train-labels.idx1-ubyte'

TEST_IMG_PATH = './MNIST_data/t10k-images-idx3-ubyte/t10k-images.idx3-ubyte'
TEST_LABEL_PATH = 'MNIST_data/t10k-labels-idx1-ubyte/t10k-labels.idx1-ubyte'


def get_img_datas(file_path):
    """
   读取MNIST数据集中的图片数据
   :param file_path: 图片文件路径
   :return: 读取到的图片文件数据列表
   """
    # 图片文件
    image_file = None
    # with open 打开文件后，出作用域，文件会自动关闭
    with open(file_path, 'rb') as file:
        image_file = file.read()

    # image数据的list，一张图片的二进制数据，一个元素
    image_datas = []
    # 文件指针，表示读取到的文件的位置
    file_index = 0
    # '>IIII'，其实>表示大端法，IIII表示读取4次Int型数据(一个Int4个字节)，依次分配到4个变量上
    # image_file，读取文件的路径
    # image_index 读取文件的开始位置
    magic_num, data_num, rows, cols = struct.unpack_from('>IIII', image_file, file_index)
    # 使文件指针向后移动4个Int大小，即跳过已经读取数据的大小
    file_index += struct.calcsize('>IIII')

    # 一次读取的大小为一张图片占用的字节数
    read_formate = ">{}B".format(rows * cols)
    for image_index in range(0, data_num):
        # 一次读取784个字节，这是因为图片是28*28大小的关系
        data = struct.unpack_from(read_formate, image_file, file_index)
        # 将一维数据，重建为二维数据，因为图片是一个矩形数据
        image_datas.append(np.reshape(data, (28, 28)))
        # 让指针向后移动784Byte
        file_index += struct.calcsize(read_formate)

    print('read image datas {}/{}'.format(len(image_datas), data_num))
    return image_datas


def get_label_datas(file_path):
    """
	读取MNIST数据集中的图片数据
	:param file_path: 标签文件路径
	:return: 读取到的标签文件数据列表
	"""
    # 标签文件
    label_file = None
    # with open 打开文件后，出作用域，文件会自动关闭
    with open(file_path, 'rb') as file:
        label_file = file.read()

    # image数据的list，一张图片的二进制数据，一个元素
    label_datas = []
    # 文件指针，表示读取到的文件的位置
    file_index = 0
    # 读取标签描述
    magic_num, data_num = struct.unpack_from('>II', label_file, file_index)
    # 使文件指针向后移动2个Int大小，即跳过已经读取数据的大小
    file_index += struct.calcsize('>II')

    for label_index in range(0, data_num):
        # 一次读取1个字节
        data = struct.unpack_from('>B', label_file, file_index)
        label_datas.append(data)
        # 让指针向后移动1个字节
        file_index += struct.calcsize('>B')

    print('read label datas {}/{}'.format(len(label_datas), data_num))
    return label_datas


# def pre_process(image_data, label_data):
# 	# 图片数据目前无需处理，转换为numpy数组，传入模型
#     image_data = np.array(image_data)
#     # 标签数据转置，与图像数据一一对应
#     label_data = np.reshape(label_data, len(label_data))
#     return  image_data, label_data
#
# train_label = get_label_datas(TRAIN_LABEL_PATH)
# train_image = get_img_datas(TRAIN_IMG_PATH)
# test_label = get_label_datas(TEST_LABEL_PATH)
# test_image = get_img_datas(TEST_IMG_PATH)
#
# # 进行预处理
# train_image,train_label = pre_process(train_image, train_label)
#
# model = models.Sequential()
# model.add(layers.Flatten(input_shape=(28,28)))
# model.add(layers.Dense(15,activation = 'relu' ))
# model.add(layers.Dense(10,activation = 'sigmoid'))
# model.compile(optimizer=tf.keras.optimizers.SGD(learning_rate=1e-3),
#               loss='mean_squared_error',
#               metrics=['accuracy'])
# model.summary()
#
# history = model.fit(x=train_image, #预测依据
#                     y=train_label, #预测结果
#                     epochs= 15)    #一共训练15次

def pre_process(image_data, label_data):
    # 转换为numpy数组
    image_data = np.array(image_data)
    # 进行归一化处理
    image_data = image_data / 255.
    # 标签数据转置，与图像数据一一对应
    label_data = np.reshape(label_data, len(label_data))
    return image_data, label_data


print(tf.version.VERSION)
print(np.version.version)
print(pd._version.version_json)

train_label = get_label_datas(TRAIN_LABEL_PATH)
train_image = get_img_datas(TRAIN_IMG_PATH)
test_label = get_label_datas(TEST_LABEL_PATH)
test_image = get_img_datas(TEST_IMG_PATH)

# 进行预处理
train_image, train_label = pre_process(train_image, train_label)
test_image, test_label = pre_process(test_image, test_label)

model = models.Sequential()
model.add(layers.Flatten(input_shape=(28, 28)))
model.add(layers.Dense(128, activation='relu'))
model.add(layers.Dropout(0.2))
model.add(layers.Dense(10, activation='softmax'))
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
model.summary()

history = model.fit(x=train_image,  # 预测依据
                    y=train_label,  # 预测结果
                    epochs=15)  # 一共训练15次

import matplotlib.pyplot as plt


def plot_metric(historyDic, metric):
    train_metrics = historyDic[metric]
    epochs = range(1, len(train_metrics) + 1)
    plt.plot(epochs, train_metrics, 'bo--')
    plt.title('Training and validation ' + metric)
    plt.xlabel("Epochs")
    plt.ylabel(metric)
    plt.legend(["train_" + metric])
    plt.savefig("".join('./' + metric + '.png'))
    plt.show()


plot_metric(history.history, "accuracy")
plot_metric(history.history, "loss")

model.save('./mnist_model.h5')
json.dump(history.history, open('history.json', 'w'))

model.evaluate(test_image, test_label, verbose=2)

model = models.Sequential()
model.add(layers.Flatten(input_shape=(28, 28)))
model.add(layers.Dense(15, activation='relu'))
model.add(layers.Dense(15, activation='relu'))
model.add(layers.Dense(10, activation='softmax'))
# 修改优化器为 adam
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])
model.summary()
